import tkinter as tk
from tkinter import ttk
import os
import subprocess
import threading
import time
import signal
from tkinter import messagebox
import traceback
from ruamel.yaml import YAML

# -------------------------- 1. 全局训练状态变量 --------------------------
_train_status = {
    "is_running": False,  # 训练进程是否在运行
    "process": None,  # 训练子进程对象
    "thread": None,  # 监控线程
    "output_log": []  # 训练输出日志
}

# -------------------------- 2. 路径配置（完全匹配目录结构） --------------------------
# 当前文件（train_operations.py）所在目录：yolo_front/
CURRENT_FILE_DIR = os.path.dirname(os.path.abspath(__file__))

# 工作目录：第一层ultralytics/（与你手动执行命令的目录一致）
# 从yolo_front/向上一级就是第一层ultralytics/
WORK_DIR = os.path.abspath(os.path.join(CURRENT_FILE_DIR, ".."))

# train.yaml的相对路径：从工作目录（第一层ultralytics/）进入第二层ultralytics/cfg/
# 对应手动命令中的 "ultralytics/cfg/train.yaml"
TRAIN_YAML_PATH = "ultralytics/cfg/train.yaml"

# 虚拟环境激活命令：从工作目录（第一层ultralytics/）出发，路径与手动操作一致
ACTIVATE_CMD = ".venv/Scripts/activate"


# -------------------------- 3. 参数设置窗口（保留原逻辑，修正参数保存路径） --------------------------
def cmd_train_params(app):
    param_window = tk.Toplevel(app)
    param_window.title("训练参数设置")
    param_window.geometry("600x400")

    # 初始化参数（默认值与train.yaml对应）
    param_vars = {
        "Epochs": tk.StringVar(),
        "cache": tk.StringVar(),
        "patience": tk.StringVar(),
        "device": tk.StringVar(),
        "batch": tk.StringVar(),
        "imgsz": tk.StringVar(),
        "save": tk.StringVar(),
        "workers": tk.StringVar(),
        "project": tk.StringVar(),
        "name": tk.StringVar()
    }

    # 使用ruamel.yaml来保留注释
    from ruamel.yaml import YAML
    yaml = YAML()
    yaml.preserve_quotes = True

    # 读取现有配置
    full_yaml_path = os.path.join(WORK_DIR, TRAIN_YAML_PATH)
    if os.path.exists(full_yaml_path):
        with open(full_yaml_path, "r", encoding="utf-8") as f:
            existing_params = yaml.load(f) or {}
    else:
        existing_params = {}

    # 设置默认值
    defaults = {
        "epochs": 300,
        "cache": "ram",
        "patience": 50,
        "device": 0,
        "batch": 16,
        "imgsz": 640,
        "save": True,
        "workers": 8,
        "project": "results",
        "name": "train"
    }

    # 合并现有参数和默认参数
    for key in param_vars:
        lower_key = key.lower()
        if lower_key in existing_params and existing_params[lower_key] is not None:
            param_vars[key].set(str(existing_params[lower_key]))
        else:
            param_vars[key].set(str(defaults.get(lower_key, "")))

    # 布局参数输入框
    for i, (param_name, var) in enumerate(param_vars.items()):
        row, col = divmod(i, 2)
        label = ttk.Label(param_window, text=param_name)
        label.grid(row=row, column=col * 2, padx=10, pady=10, sticky="w")
        entry = ttk.Entry(param_window, textvariable=var)
        entry.grid(row=row, column=col * 2 + 1, padx=10, pady=10)

    # 占位标签
    ttk.Label(param_window, text="……", font=("Arial", 20)).grid(
        row=0, column=4, rowspan=5, padx=20, pady=20
    )

    # 确认按钮：保存参数到正确的train.yaml路径
    def confirm_params():
        try:
            # 重新加载配置以保留注释
            full_yaml_path = os.path.join(WORK_DIR, TRAIN_YAML_PATH)
            if os.path.exists(full_yaml_path):
                with open(full_yaml_path, "r", encoding="utf-8") as f:
                    config = yaml.load(f) or {}
            else:
                config = {}

            # 更新参数值（保留注释和其他未更改的配置）
            # 对特定参数进行类型转换
            for name, var in param_vars.items():
                key = name.lower()
                value = var.get()
                
                # 对数值类型参数进行转换
                if key in ["epochs", "patience", "batch", "imgsz", "workers"]:
                    try:
                        config[key] = int(value)
                    except ValueError:
                        config[key] = defaults[key]  # 使用默认值
                elif key == "save":
                    # 布尔类型转换
                    if value.lower() in ['true', '1', 'yes', 'on']:
                        config[key] = True
                    elif value.lower() in ['false', '0', 'no', 'off']:
                        config[key] = False
                    else:
                        config[key] = defaults[key]  # 使用默认值
                elif key == "device":
                    # device 可能是整数或字符串
                    try:
                        config[key] = int(value)
                    except ValueError:
                        config[key] = value  # 保留为字符串，如 "cpu"
                else:
                    # 其他参数保持为字符串
                    config[key] = value

            # 写回文件（保留注释）
            with open(full_yaml_path, "w", encoding="utf-8") as f:
                yaml.dump(config, f)

            app.show_status(f"参数已保存到: {full_yaml_path}")
            param_window.destroy()
        except Exception as e:
            messagebox.showerror("保存失败", f"错误信息: {str(e)}")

    confirm_button = ttk.Button(param_window, text="确认", command=confirm_params)
    confirm_button.grid(row=10, column=0, columnspan=4, pady=20)


# -------------------------- 4. 训练日志读取（显示所有日志，方便排查错误） --------------------------
def _read_process_output(process, app):
    while process.poll() is None and _train_status["is_running"]:
        try:
            line = process.stdout.readline()
            if line:
                # 解码日志（兼容特殊字符）
                line = line.decode("utf-8", errors="replace").strip()
                _train_status["output_log"].append(line)
                # 实时显示所有日志（不过滤，方便定位问题）
                app.show_status(line)
                # 限制日志长度，避免内存占用过大
                if len(_train_status["output_log"]) > 2000:
                    _train_status["output_log"].pop(0)
        except Exception as e:
            app.show_status(f"日志读取错误: {str(e)}")
            break

    # 训练结束反馈
    exit_code = process.poll()
    if exit_code == 0:
        app.show_status("✅ 训练完成！结果已保存到 project/name 目录")
    else:
        app.show_status(f"❌ 训练异常终止，退出码: {exit_code}（查看日志找具体原因）")

    # 重置训练状态
    _train_status["is_running"] = False
    _train_status["process"] = None
    _train_status["thread"] = None


# -------------------------- 5. 启动训练（核心：与手动命令100%一致） --------------------------
def cmd_train_start(app):
    # 检查是否已在训练
    if _train_status["is_running"]:
        app.show_status("⚠️  训练已在运行中，无需重复启动")
        return

    # 先验证train.yaml是否存在（避免路径错误）
    full_yaml_path = os.path.join(WORK_DIR, TRAIN_YAML_PATH)
    if not os.path.exists(full_yaml_path):
        messagebox.showerror("文件缺失", f"未找到train.yaml：\n{full_yaml_path}")
        return

    # 构建PowerShell命令（完全模仿手动操作）
    # 命令逻辑：激活虚拟环境 → 执行yolo训练（参数与手动一致）
    cmd = [
        "powershell.exe", "-Command",
        f"& '{ACTIVATE_CMD}'; yolo cfg={TRAIN_YAML_PATH}"
    ]

    try:
        # 启动子进程（工作目录设置为第一层ultralytics/）
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,  # 合并 stderr 到 stdout，方便查看错误
            shell=False,
            cwd=WORK_DIR,  # 关键：与手动执行命令的目录完全一致
            preexec_fn=os.setsid if os.name != "nt" else None
        )

        # 更新训练状态
        _train_status["is_running"] = True
        _train_status["process"] = process
        _train_status["output_log"] = []

        # 启动日志监控线程（不阻塞UI）
        _train_status["thread"] = threading.Thread(
            target=_read_process_output,
            args=(process, app),
            daemon=True  # 主线程退出时自动关闭该线程
        )
        _train_status["thread"].start()

        app.show_status(f"🚀 训练已启动！工作目录：{WORK_DIR}")
        app.show_status(f"📜 执行命令：powershell -Command {ACTIVATE_CMD}; yolo cfg={TRAIN_YAML_PATH}")

    except Exception as e:
        app.show_status(f"❌ 启动训练失败：{str(e)}")
        messagebox.showerror("启动错误", f"详细信息：\n{traceback.format_exc()}")


# -------------------------- 6. 停止训练（安全终止进程） --------------------------
def cmd_train_stop(app):
    if not _train_status["is_running"] or not _train_status["process"]:
        app.show_status("⚠️  没有运行中的训练进程")
        return

    try:
        process = _train_status["process"]
        app.show_status("🛑 正在终止训练进程...")

        # Windows系统：先发送Ctrl+C（优雅终止，保留已训练结果），失败则强制终止
        if os.name == "nt":
            try:
                os.kill(process.pid, signal.CTRL_C_EVENT)
            except:
                os.kill(process.pid, signal.SIGTERM)
        # Linux/macOS（备用，若你后续切换系统）
        else:
            os.killpg(os.getpgid(process.pid), signal.SIGTERM)

        # 等待进程终止（最多等5秒）
        timeout = 5
        while process.poll() is None and timeout > 0:
            time.sleep(0.5)
            timeout -= 0.5

        # 强制终止（若5秒后仍未退出）
        if process.poll() is None:
            if os.name == "nt":
                os.kill(process.pid, signal.SIGKILL)
            else:
                os.killpg(os.getpgid(process.pid), signal.SIGKILL)

        # 重置状态
        _train_status["is_running"] = False
        _train_status["process"] = None
        _train_status["thread"] = None
        app.show_status("✅ 训练进程已终止")

    except Exception as e:
        app.show_status(f"❌ 停止训练失败：{str(e)}")
        messagebox.showerror("停止错误", f"详细信息：\n{traceback.format_exc()}")


# -------------------------- 7. 训练状态查询（供UI显示） --------------------------
def get_train_status():
    return {
        "is_running": _train_status["is_running"],
        "log_count": len(_train_status["output_log"]),
        "last_log": _train_status["output_log"][-1] if _train_status["output_log"] else "无日志"
    }


# -------------------------- 8. 训练日志获取（供UI展示完整日志） --------------------------
def get_train_logs(max_lines=200):
    return _train_status["output_log"][-max_lines:]  # 返回最后200行日志